#!/usr/local/bin/python
# Rob Dennis - Lurchington Q. Delacroix
# Keasby - TV renamer
# Project Homepage: http://code.google.com/p/keasby/

#--- Python imports
import os
import re
import sys
import logging
import urllib2
import random
import shelve
import anydbm
import time
from os.path import join as pjoin
from os.path import exists as pexists
#--- Keasby Imports
from keasby_repo import KeasbySource, KeasbySourceTestCase
from keasby_cache import KeasbyCache, KeasbyCacheError
from keasby_toolbox import all_keys, dump, choose_single, swap_keys
from xml_tools import XmlHandler
from tv_episode import TvEpisode

random.seed()

class TvDbApi(KeasbySource):
    """
    This class is written for interacting with thetvdb.com's API
    """

    def __init__(self,api_key):
        """
        Instantiate's class attributes for caching, interaction with data,
        and for the initialization of mirrors and updates since last
        interaction.
        """

        self.cache = TvDbCache("http://www.thetvdb.com", api_key)
        self.data = TvDbData(self.cache.dir)
        self.xml = XmlHandler(self.cache.dir)
        self.tvepisode_mapping_dict = dict(episode_number = "EpisodeNumber",
                                            season_number = "SeasonNumber",
                                            title = "EpisodeName",
                                            showname = "SeriesName",
                                            seriesid = "id")

        self.search_type_map = dict(season = "season",
                                    absolute = "absolute",
                                    dvd = "dvd")
        self._init_mirrors()
        self._init_language()
        self._init_current_info()

        KeasbySource.__init__(self, "TVDB",
                                "API for interaction with www.thetvdb.com")

#---Initialization

    def _init_mirrors(self):
        """
        Downloads the available mirrors.xml file, and then populates
        the mirrors as xml, banner, or zip, or a combination of any or
        all
        """

        if not "mirrors.xml" in self.cache:
            mirror_xml_dump = dump(self.cache.api_cache("mirrors.xml"))
            self.xml.parse_xml_string("mirrors.xml", mirror_xml_dump)

        else:
            self.xml.parse_xml_path(pjoin(self.cache.dir, "mirrors.xml"))

        self.mirror = {"xml":[],"banner":[],"zip":[]}

        for mirror in self.xml["mirrors.xml"]["Mirror"]:
            typemask = int(mirror["typemask"][0].data())
            mirrorpath = mirror["mirrorpath"][0].data()

            if typemask & 1 == 1:
                self.mirror["xml"].append(mirrorpath.encode("UTF-8"))
            if typemask & 2 == 2:
                self.mirror["banner"].append(mirrorpath.encode("UTF-8"))
            if typemask & 4 == 4:
                self.mirror["zip"].append(mirrorpath.encode("UTF-8"))

    def _init_language(self):
        """
        Currently defaults to English, but could change
        """

        self.language = u"en"

    def _init_current_info(self):
        """
        Downloads and stores the most recent updates.xml
        """

        updates_dump = self.cache.api_cache("updates/",
                                            filename="updates.xml",
                                            overwrite=True)

        _utime = self.xml["updates.xml"]["Data"][0].attributes["time"].value

        if _utime.isdigit():
            self.time = time.gmtime(float(_utime))
        else:
            raise TvDbApiError("unknown error getting last update time")

#---Source Interface

    def query_search(self, kwargs):
        """
        Checks to see if showname has alredy been mapped to a series ID.
        If so, return it, otherwise query the showname and return the possible
        choices
        """

        showname = self._cast_kwargs(kwargs).get("showname")

        if showname not in self.data:
            #it's not already loaded into data
            return self._get_seriesid_list(showname)

        else:
            return {"seriesid":self.data[showname]}


    def downselect(self, choice, kwargs):
        """
        takes in the selection, and maps the seriesid to the provided
        showname value
        """

        kwargs = self._cast_kwargs(kwargs)

        showname = kwargs.get("showname")
        seriesid = choice["seriesid"]

        if not showname:
            raise TvDbApiError("invalid showname: %s"%showname)
        if not isinstance(showname,basestring):
            raise TvDbApiError("showname must be of type string or unicode"+
                                " %s"%type(showname))
        else:
            #setting the seriesid
            self.data.set_id(showname, seriesid)

    def prime_search(self, kwargs):
        """
        takes in the search arguments and uses the mapped seriesid to download,
        parse, and arrange the series XML file.
        """

        kwargs = self._cast_kwargs(kwargs)

        showname = kwargs.get("showname")
        seriesid = self.data.get_id(showname)

        if not showname:
            raise TvDbApiError("invalid showname: %s"%showname)
        if not isinstance(showname,basestring):
            raise TvDbApiError("showname must be of type string or unicode"+
                                " %s"%type(showname))
        if not seriesid:
            raise TvDbApiError("invalid seriesid: %s"%seriesid)

        #parsing the data for showname

        inferred_search = self._infer_search(kwargs)

        if inferred_search is self.search_type_map["season"]:
            prepped_data = self._season_prep(self._map_series(seriesid))

        elif inferred_search is self.search_type_map["absolute"]:
            prepped_data = self._absolute_prep(self._map_series(seriesid))

        elif inferred_search is self.search_type_map["dvd"]:
            prepped_data = self._dvd_prep(self._map_series(seriesid))

        else:
            raise TvDbApiError("invalid search terms supplied: %s"%kwargs)

        if seriesid not in self.data:
            logging.debug("initializing showname: %s", showname)
            self.data[seriesid] = {}

        #hack: bracket notation isn't working quite right
        existing_dict = self.data[seriesid].copy()
        existing_dict.update({inferred_search:prepped_data})
        self.data[seriesid] = existing_dict

    def get_search(self, kwargs):
        """
        returns the search mapping, modified to include all the
        information found by TvDbApi
        """

        found_args = self._conduct_search(self._cast_kwargs(kwargs))

        return found_args

    def _conduct_search(self, kwargs):
        """
        conduct the search after ensuring that the information is there
        """

        kwargs = kwargs.copy()

        if "seriesid" not in kwargs:
            logging.debug("loading seriesid map from data")

            if "showname" not in kwargs:
                raise(TvDbApiError,"incorrect search terms supplied: %s"%kwargs)

            kwargs["seriesid"] = unicode(self.data.get_id(kwargs["showname"]))

        inferred_search = self._infer_search(kwargs)

        if inferred_search is self.search_type_map["season"]:
            return self._search_season(kwargs)

        elif inferred_search is self.search_type_map["absolute"]:
            return self._search_absolute(kwargs)

        elif inferred_search is self.search_type_map["dvd"]:
            return self._search_dvd(kwargs)

        else:
            raise(TvDbApiError,"incorrect search terms supplied: %s"%kwargs)

#---Search Backend

    def _search_season(self, kwargs):
        """
        Searches for title based on TVDB's season ordering
        """

        try:
            seriesid, season_number, episode_number = all_keys(kwargs,
                                                                "seriesid",
                                                                "season_number",
                                                                "episode_number")
        except Exception, error_message:
            raise TvDbApiError(error_message)


        logging.debug("accessing data for %s with:", self.__class__.__name__)
        logging.debug("seriesid %s", seriesid)
        logging.debug("season number %s", season_number)
        logging.debug("episode number %s", episode_number)



        try:
            return self.data[seriesid]["season"][season_number][episode_number]
        except KeyError, error_message:
            raise TvDbApiError("%s in: %s"%(error_message, self.data[showname]))

    def _search_absolute(self, kwargs):
        """
        Searches for title based on TVDB's absolute ordering
        """
        #abs_args = all_keys(kwargs, "showname", "absolute_number")
        try:
            seriesid, absolute_number = all_keys(kwargs, "seriesid",
                                                         "absolute_number")
        except Exception, error_message:
            raise TvDbApiError(error_message)

        raise NotImplementedError("Absolute ordering not yet implemented in %s"%
                                self.__class__.__name__)

    def _search_dvd(self, kwargs):
        """
        Searches for title based on DVD ordering
        """
        #dvd_args = all_keys(kwargs, "showname", "season_number",
        #                    "dvd_disc", "disc_episode")
        try:
            seriesid, season_number, dvd_disc, disc_episode = all_keys(kwargs,
                                                            "seriesid",
                                                            "season_number",
                                                            "dvd_disc",
                                                            "disc_episode")
        except Exception, error_message:
            raise TvDbApiError(error_message)

        raise NotImplementedError("DVD ordering not yet implemented in %s"%
                                    self.__class__.__name__)

#---Prepping the data for later interaction
##
##    def _season_order(self, series_map):
##        """
##        Takes in an array of results dictionaries and splits them based on
##        the season_number into keys of a dictionary, attempts to
##        order them by unique keys
##        """
##
##        return self._dict_tree(series_map, ["season_number", "episode_number"])
##
##    def _dict_tree(self, series_map, key_list):
##        """
##        Takes in an array of dictionaries and creates a number of sub-
##        dictionaries based on each key in key_list.
##
##        For example, with a key_list of ["season_number", "episode_number"]
##        I'd be expecting to get access to a particular mapping by getting:
##        [season_number][episode_number]
##
##        """
##
##        if not key_list:
##            raise TvDbApiError("empty key_list: %s"%key_list)
##
##        if not series_map:
##            raise TvDbApiError("empty series_map: %s"%series_map)
##
##        tldict = {}
##
##        for series in series_map:
##            #dictionary assignment is actually a pointer
##            opdict = tldict
##
##
##            if len(key_list) > 1:
##                #creates all the tiers based on the order of keys,
##                #except for the last
##                for key in key_list[:-1]:
##                    series_value = series[key]
##                    if series_value not in opdict:
##                        opdict[series_value] = {}
##                    opdict = opdict[series_value]
##
##            #the last key in the hierarchy actually holds the map
##            opdict[series[key_list[-1]]] = series
##
##
##        return tldict
##
###---Interaction with TvDbData class

    def _get_seriesid_list(self, showname):
        """
        Get all the available series listing from a given showname
        """

        return self._parse_showname_query(showname)

    def _select_seriesid(self, showname):
        """
        Select the seriesid to tie with a given showname search
        """

        chosen_entry = choose_single(self._parse_showname_query(showname))

        return chosen_entry["seriesid"]

#---XML interaction

    def _map_series(self, seriesid):
        """
        shapes the output of series as an array of TvEpisodes
        """

        header,records = self._parse_series_record(seriesid)
        header_dict = header[0]

        series_map = []

        for record in records:
            #the header of the xml file holds series static information
            #like showname, series id, etc, that record does not have
            temp_dict = header_dict.copy()
            temp_dict.update(record)

            series_map.append(swap_keys(temp_dict,self.tvepisode_mapping_dict))

        return series_map


    def _parse_showname_query(self, showname):
        """
        Shortcut for converting the XML-DOM returned from the GetSeries query
        """

        item_keys = [u"seriesid",u"SeriesName"]
        optional_keys = [u"FirstAired",u"Overview"]

        query_content = self._query_showname(showname)
        self.xml.parse_xml_string("%s.search"%showname,query_content)

        return self._xml_array("%s.search"%showname,
                                u"Series",
                                item_keys,
                                optional_keys)

    def _parse_series_record(self, seriesid):
        """
        Shortcut for converting the XML-DOM contained in the seriesid zipfile
        Returns the showname and episode arrays
        """

        series_keys = [u"id",u"SeriesName",u"lastupdated"]

        ep_keys = [u"EpisodeName",u"EpisodeNumber",
                    u"SeasonNumber",u"absolute_number",u"lastupdated"]
        opt_keys = [u"FirstAired"]

        try:
            self.xml.parse_zipfile("%s.id"%seriesid,
                                    self._cache_series_zip(seriesid),
                                    "%s.xml"%self.language)

        except Exception, error_message:

            raise TvDbCacheError("Error parsing the series record zip"+
                                " seriesid: %s"%seriesid+
                                " path: %s"%self._cache_series_zip(seriesid)+
                                " language: %s.xml"%self.language+
                                " : %s"%error_message)

        return (self._xml_array("%s.id"%seriesid,u"Series",series_keys),
                self._xml_array("%s.id"%seriesid,u"Episode",ep_keys,opt_keys))


    def _xml_array(self, xml_key, array_key, item_keys, optional_keys=[]):
        """
        Returns the dictionary array from XML-DOM,
        """

        return self.xml[xml_key].get_array(array_key,item_keys,optional_keys)

#---TvDbData interaction

    def _season_prep(self, series_map):
        """
        Orders by season_number, then episode_number
        """

        return self._dict_tree(series_map,["season_number", "episode_number"])

    def _dvd_prep(self, series_map):
        """
        Orders by season_number, then dvd_disc, then disc_episode
        """

        return self._dict_tree(series_map,["season_number",
                                            "dvd_disc",
                                            "disc_episode"])

    def _absolute_prep(self, series_map):
        """
        Orders by absolute_number
        """

        return self._dict_tree(series_map,["absolute_number"])

    def _dict_tree(self, series_map, key_list):
        """
        Takes in an array of dictionaries and creates a number of sub-
        dictionaries based on each key in key_list.

        For example, with a key_list of ["season_number", "episode_number"]
        I'd be expecting to get access to a particular mapping by getting:
        [season_number][episode_number]

        """

        if not key_list:
            raise TvDbApiError("empty key_list: %s"%key_list)

        if not series_map:
            raise TvDbApiError("empty series_map: %s"%series_map)

        tldict = {}

        for series in series_map:
            #dictionary assignment is actually a pointer
            opdict = tldict


            if len(key_list) > 1:
                #creates all the tiers based on the order of keys,
                #except for the last
                for key in key_list[:-1]:
                    series_value = series[key]
                    if series_value not in opdict:
                        opdict[series_value] = {}
                    opdict = opdict[series_value]

            #the last key in the hierarchy actually holds the map
            opdict[series[key_list[-1]]] = series


        return tldict

#---URL Interaction

    def _query_showname(self, showname):
        """
        Queries the tvdb api based on showname.
        Returns the XML data that will need to be parsed out
        """

        _showname = self._scrub_showname(showname)
        _escaped_showname = self._escape_text(_showname)

        return self.cache.query_cache("GetSeries.php?seriesname=%s"%
                                                            _escaped_showname)

    def _cache_series_zip(self, seriesid):
        """
        Downloads the all series zip file and returns the contents
        """

        seriesid_path = ("series",seriesid,"all","%s.zip"%self.language)

        url = "/".join(seriesid_path)
        cache_path = pjoin(*seriesid_path[:-1])

        if cache_path in self.cache:
            logging.debug("loading pre-existing information for %s",seriesid)

        else:
            logging.debug("downloading information for %s",seriesid)
            logging.debug("api caching with path: %s",cache_path)

        return self.cache.api_cache(url,self._mirror("zip"),cache_path)



    def _mirror(self, mirror_type):
        """
        Return a mirror url of the appropriate type, suitable for interaction
        with TvDbCache
        """

        try:
            random_index = random.randint(0,len(self.mirror[mirror_type])-1)
        except Exception,error_message:
            raise TvDbApiError("unable to select a %s-mirror"%mirror_type)

        return self.mirror[mirror_type][random_index]

    def _scrub_showname(self, showname):
        """
        remove punctuation and other characters that won't be searched
        """

        return showname

    def _escape_text(self, text):

        return re.sub(r"\s+",r"%20",text)



#---Utility

    def _test_search_key(self, search_key, search_dict):
        """
        Convenience class for ensuring that a key exists in
        the provided search dictionary
        """

        if search_key in search_dict:
            return search_dict[search_key]
        else:
            raise TvDbApiError("required search term %s not given"%search_key)

    def _infer_search(self, kwargs):
        """
        Based on the keys available in kwargs, return which search type
        is appropriate
        """

        if all_keys(kwargs, "showname", "season_number", "episode_number"):
            return self.search_type_map["season"]

        if all_keys(kwargs, "showname", "absolute_number"):
            return self.search_type_map["absolute"]

        if all_keys(kwargs, "showname",
                    "season_number", "dvd_disc", "disc_episode"):
            return self.search_type_map["dvd"]

    def _cast_kwargs(self, kwargs):
        """
        casts the input key-word arguments into the necessary format for
        use in the TvDbApi
        """

        for key, value in kwargs.items():
            if not isinstance(value, unicode):
                kwargs[key] = unicode(value)

        return kwargs


class TvDbData():
    def __init__(self, data_path=""):
        """
        Class for interacting with TvDb data
        """

        self.data_path = data_path
        self.open()

#---Initialization / Cleanup

    def open(self):
        """
        Opens all data resources for later interaction
        """

        self.data = shelve.open(pjoin(self.data_path,"tvdb_data.cache"),"c")
        self.seriesid = anydbm.open(pjoin(self.data_path,"tvdb_id.map"),"c")

    def close(self):
        """
        Closes open data sources in order to cleanly release of resources
        """

        self.data.close()
        self.seriesid.close()

    def sync(self):
        """
        Forces a sync after modification of values
        """

        self.data.sync()

#---Interaction with Data

    def set_data(self, showname, value):
        """
        Shelves the value to key: <showname>, then syncs the shelf
        """

        self.data[self._cast_string(showname)] = value
        self.data.sync()


    def get_data(self, showname):
        """
        Returns the the value associated with showname. Can throw a
        KeyError on invalid showname
        """

        return self.data[self._cast_string(showname)]

#---Interaction with Series ID Numbers

    def get_id(self, showname):
        """
        Returns seriesid for a supplied showname. Returns None
        if the showname is not currently mapped
        """

        return self.seriesid.get(self._cast_string(showname))

    def set_id(self, showname, seriesid):
        """
        Sets the supplied showname as the given seriesid
        """

        self.seriesid[self._cast_string(showname)] = seriesid

#---Operator

    def __contains__(self, showname):
        """
        overloads in operator for testing the presence of a showname
        """

        return self._cast_string(showname) in self.data

    def __getitem__(self, showname):
        """
        Currently overloads get_data
        """

        return self.get_data(self._cast_string(showname))

    def __setitem__(self, showname, value):
        """
        Currently overloads set_data
        """

        return self.set_data(showname,value)

#---Utility

    def _cast_string(self, text):
        """
        used to ensure that unicode is encoded as a string for interaction
        with shelve and anydbm which require it
        """

        if isinstance(text,unicode):
            return text.encode("UTF-8")
        else:
            return text


class TvDbCache(KeasbyCache):
    """
    Subclass of KeasbyCache that intends to create shortcuts for caching
    information from thetvdb.com
    """

    def __init__(self, base_url, api_key):
        """
        Class for caching URL files to disk, and the
        information  to plain text.
        """

        self.base_url = base_url
        self.api_key = api_key

        KeasbyCache.__init__(self)

#---Cache Interface

    def api_cache(self, url, netloc=None, cache_dir="", filename=None,
                    overwrite=False, write=True):
        """
        Caches the contents of the relative url to netloc, inserts
        /api/<apikey>/ and returns the string that was cached
        """

        api_url = self._get_api_url(url)
        pass_through_list = [cache_dir,filename,overwrite,write]

        return self.url_cache(api_url,*pass_through_list)

    def query_cache(self, url, cache_dir="", filename=None,
                        overwrite=False, write=False):
        """
        Caches the contents of the relative url to netloc, and returns
        the string that was cached
        """

        query_url = self._make_url(self.base_url,"api",url)
        pass_through_list = [cache_dir,filename,overwrite,write]

        return self.url_cache(query_url,*pass_through_list)

    def zip_cached(self, seriesid, language=u"en"):
        """
        Convenience functions for testing if a seriesid/language.zip
        file
        """

        if seriesid is not None:
            return pjoin("series",seriesid,"all","%s.zip"%language) in self
        else:
            return False

    def _shape_cache_info(self, cache_dir, filename,url):
        """
        Convenience function for shaping the cache and filename
        """

        if filename is None:
            filename = self._get_url_filename(url)
        cache_dir = self._get_cache_path(cache_dir)

        return cache_dir,filename

#---Utility

    def _get_api_url(self, url_path, url_netloc=None):
        """
        Convenience function for interacting accessing a URL relative to
        the tvdb API key
        """

        if url_netloc is None:
            url_netloc = self.base_url


        return self._make_url(url_netloc,"api",self.api_key,url_path)



class TvDbApiError(Exception):
    """
    Generic TvDbApi Exception
    """
    pass

class TvDbCacheError(TvDbApiError,KeasbyCacheError):
    """
    Generic TvDbCache Exception
    """
    pass

class TvDbDataError(TvDbApiError):
    """
    Generic TvDbDataError
    """
    pass


class TvDbApiTestCase(KeasbySourceTestCase):
    """
    Tests the use of the TvDbApi class
    """

    def setUp(self):
        """
        Instantiates a KeasbyOptions instance
        """

        self.s = TvDbApi("D1EBA6781E2572BB")

    def tearDown(self):
        """
        Closes the cache and data attributes of the TvDbApi class
        in order to ensure a clear environment for future tests.
        """

        self.s.data.close()
        self.s.cache.purge()

    def test_sanity(self):
        """
        basic test here
        """

        self.assertLog(self.s)
        self.assertLog(pexists(self.s.cache.dir),"cache directory wasn't made")

    def test_internal_initialization(self):
        """
        tests for the proper initialization of class attributes and resources
        """

        self.assertLog(self.s.xml["mirrors.xml"])
        test_url = r"http://thetvdb.com"

        for mirror_type, mirror_list in self.s.mirror.iteritems():
            self.assertLog(test_url in mirror_list,
                            "%s not in %s list"%(test_url,mirror_type))

            self.assertLog(isinstance(self.s.time,time.struct_time))

    def test_internal_xml_parsing(self):
        """
        tests the ability of the XML handler to parse TVDB-specific data
        """

        test_dict = {"SeriesName":u"Lost", "seriesid":u"73739"}
        episode_dict = {"EpisodeNumber":u"1", "SeasonNumber":u"1"}

        #!Hack: this will fail if ABC's Lost isn't the first result
        show_dict = self.s._parse_showname_query(test_dict["SeriesName"])[0]

        for test_key,test_value in test_dict.iteritems():
            self.assertEquals(test_value, show_dict[test_key])

        header, ep_array = self.s._parse_series_record(test_dict["seriesid"])

        self.assertLog(header, "invalid header", header)

        self.assertLog(ep_array, "invalid episode array")

        for episode in ep_array:
            self.assertLog(episode, "invalid episode", episode)

    def test_internal_dict_tree(self):
        """
        tests the ability to construct a "tree" based on a dict and
        an key_list.
        """

        map_array = [{"number":"1", "letter":"a"},
                        {"number":"1", "letter":"b"},
                        {"number":"1", "letter":"c"},
                        {"number":"2", "letter":"a"},
                        {"number":"2", "letter":"b"},
                        {"number":"2", "letter":"c"},
                        {"number":"3", "letter":"a"},
                        {"number":"3", "letter":"b"},
                        {"number":"3", "letter":"c"}]

        expect_dict = {}
        expect_dict = expect_dict.fromkeys("123", {})

        self.assertLog(expect_dict,
                        "invalid dict: %s"%expect_dict,
                        expect_dict)

        for key in expect_dict.iterkeys():
            expect_dict[key] = expect_dict[key].fromkeys("abc", {})
            self.assertLog(expect_dict[key],
                            "invalid dict: %s"%expect_dict[key],
                            expect_dict[key])

        expect_dict["1"]["a"] = map_array[0]
        expect_dict["1"]["b"] = map_array[1]
        expect_dict["1"]["c"] = map_array[2]
        expect_dict["2"]["a"] = map_array[3]
        expect_dict["2"]["b"] = map_array[4]
        expect_dict["2"]["c"] = map_array[5]
        expect_dict["3"]["a"] = map_array[6]
        expect_dict["3"]["b"] = map_array[7]
        expect_dict["3"]["c"] = map_array[8]

        for map in map_array:
            self.assertLog(map, "invalid map: %s"%map, map)

        for key, value in expect_dict.iteritems():
            self.assertLog(key, "invalid key: %s"%key, key)
            self.assertLog(value, "invalid value: %s"%value, value)

        self.assertLog(expect_dict,
                        "invalid dict: %s"%expect_dict,
                        expect_dict)

        result_map = self.s._dict_tree(map_array, ["number", "letter"])

        self.assertDictEquals(result_map, expect_dict)

